{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "自然语言推理数据集.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "g4UN2a_q1g9s",
        "outputId": "26625c53-2ada-4332-e4c6-47847197ba29"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Collecting d2l\n",
            "  Downloading d2l-0.17.3-py3-none-any.whl (82 kB)\n",
            "\u001b[K     |████████████████████████████████| 82 kB 323 kB/s \n",
            "\u001b[?25hRequirement already satisfied: jupyter==1.0.0 in /usr/local/lib/python3.7/dist-packages (from d2l) (1.0.0)\n",
            "Collecting numpy==1.18.5\n",
            "  Downloading numpy-1.18.5-cp37-cp37m-manylinux1_x86_64.whl (20.1 MB)\n",
            "\u001b[K     |████████████████████████████████| 20.1 MB 79.6 MB/s \n",
            "\u001b[?25hCollecting requests==2.25.1\n",
            "  Downloading requests-2.25.1-py2.py3-none-any.whl (61 kB)\n",
            "\u001b[K     |████████████████████████████████| 61 kB 8.8 MB/s \n",
            "\u001b[?25hCollecting pandas==1.2.2\n",
            "  Downloading pandas-1.2.2-cp37-cp37m-manylinux1_x86_64.whl (9.9 MB)\n",
            "\u001b[K     |████████████████████████████████| 9.9 MB 20.8 MB/s \n",
            "\u001b[?25hCollecting matplotlib==3.3.3\n",
            "  Downloading matplotlib-3.3.3-cp37-cp37m-manylinux1_x86_64.whl (11.6 MB)\n",
            "\u001b[K     |████████████████████████████████| 11.6 MB 41.9 MB/s \n",
            "\u001b[?25hRequirement already satisfied: qtconsole in /usr/local/lib/python3.7/dist-packages (from jupyter==1.0.0->d2l) (5.2.2)\n",
            "Requirement already satisfied: notebook in /usr/local/lib/python3.7/dist-packages (from jupyter==1.0.0->d2l) (5.3.1)\n",
            "Requirement already satisfied: nbconvert in /usr/local/lib/python3.7/dist-packages (from jupyter==1.0.0->d2l) (5.6.1)\n",
            "Requirement already satisfied: ipywidgets in /usr/local/lib/python3.7/dist-packages (from jupyter==1.0.0->d2l) (7.6.5)\n",
            "Requirement already satisfied: jupyter-console in /usr/local/lib/python3.7/dist-packages (from jupyter==1.0.0->d2l) (5.2.0)\n",
            "Requirement already satisfied: ipykernel in /usr/local/lib/python3.7/dist-packages (from jupyter==1.0.0->d2l) (4.10.1)\n",
            "Requirement already satisfied: kiwisolver>=1.0.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.3->d2l) (1.3.2)\n",
            "Requirement already satisfied: python-dateutil>=2.1 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.3->d2l) (2.8.2)\n",
            "Requirement already satisfied: pyparsing!=2.0.4,!=2.1.2,!=2.1.6,>=2.0.3 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.3->d2l) (3.0.7)\n",
            "Requirement already satisfied: cycler>=0.10 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.3->d2l) (0.11.0)\n",
            "Requirement already satisfied: pillow>=6.2.0 in /usr/local/lib/python3.7/dist-packages (from matplotlib==3.3.3->d2l) (7.1.2)\n",
            "Requirement already satisfied: pytz>=2017.3 in /usr/local/lib/python3.7/dist-packages (from pandas==1.2.2->d2l) (2018.9)\n",
            "Requirement already satisfied: chardet<5,>=3.0.2 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1->d2l) (3.0.4)\n",
            "Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1->d2l) (1.24.3)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1->d2l) (2021.10.8)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.7/dist-packages (from requests==2.25.1->d2l) (2.10)\n",
            "Requirement already satisfied: six>=1.5 in /usr/local/lib/python3.7/dist-packages (from python-dateutil>=2.1->matplotlib==3.3.3->d2l) (1.15.0)\n",
            "Requirement already satisfied: tornado>=4.0 in /usr/local/lib/python3.7/dist-packages (from ipykernel->jupyter==1.0.0->d2l) (5.1.1)\n",
            "Requirement already satisfied: traitlets>=4.1.0 in /usr/local/lib/python3.7/dist-packages (from ipykernel->jupyter==1.0.0->d2l) (5.1.1)\n",
            "Requirement already satisfied: jupyter-client in /usr/local/lib/python3.7/dist-packages (from ipykernel->jupyter==1.0.0->d2l) (5.3.5)\n",
            "Requirement already satisfied: ipython>=4.0.0 in /usr/local/lib/python3.7/dist-packages (from ipykernel->jupyter==1.0.0->d2l) (5.5.0)\n",
            "Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.7/dist-packages (from ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (57.4.0)\n",
            "Requirement already satisfied: prompt-toolkit<2.0.0,>=1.0.4 in /usr/local/lib/python3.7/dist-packages (from ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (1.0.18)\n",
            "Requirement already satisfied: decorator in /usr/local/lib/python3.7/dist-packages (from ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (4.4.2)\n",
            "Requirement already satisfied: pexpect in /usr/local/lib/python3.7/dist-packages (from ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (4.8.0)\n",
            "Requirement already satisfied: pygments in /usr/local/lib/python3.7/dist-packages (from ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (2.6.1)\n",
            "Requirement already satisfied: simplegeneric>0.8 in /usr/local/lib/python3.7/dist-packages (from ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (0.8.1)\n",
            "Requirement already satisfied: pickleshare in /usr/local/lib/python3.7/dist-packages (from ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (0.7.5)\n",
            "Requirement already satisfied: wcwidth in /usr/local/lib/python3.7/dist-packages (from prompt-toolkit<2.0.0,>=1.0.4->ipython>=4.0.0->ipykernel->jupyter==1.0.0->d2l) (0.2.5)\n",
            "Requirement already satisfied: widgetsnbextension~=3.5.0 in /usr/local/lib/python3.7/dist-packages (from ipywidgets->jupyter==1.0.0->d2l) (3.5.2)\n",
            "Requirement already satisfied: ipython-genutils~=0.2.0 in /usr/local/lib/python3.7/dist-packages (from ipywidgets->jupyter==1.0.0->d2l) (0.2.0)\n",
            "Requirement already satisfied: jupyterlab-widgets>=1.0.0 in /usr/local/lib/python3.7/dist-packages (from ipywidgets->jupyter==1.0.0->d2l) (1.0.2)\n",
            "Requirement already satisfied: nbformat>=4.2.0 in /usr/local/lib/python3.7/dist-packages (from ipywidgets->jupyter==1.0.0->d2l) (5.1.3)\n",
            "Requirement already satisfied: jupyter-core in /usr/local/lib/python3.7/dist-packages (from nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (4.9.1)\n",
            "Requirement already satisfied: jsonschema!=2.5.0,>=2.4 in /usr/local/lib/python3.7/dist-packages (from nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (4.3.3)\n",
            "Requirement already satisfied: importlib-resources>=1.4.0 in /usr/local/lib/python3.7/dist-packages (from jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (5.4.0)\n",
            "Requirement already satisfied: pyrsistent!=0.17.0,!=0.17.1,!=0.17.2,>=0.14.0 in /usr/local/lib/python3.7/dist-packages (from jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (0.18.1)\n",
            "Requirement already satisfied: attrs>=17.4.0 in /usr/local/lib/python3.7/dist-packages (from jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (21.4.0)\n",
            "Requirement already satisfied: typing-extensions in /usr/local/lib/python3.7/dist-packages (from jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (3.10.0.2)\n",
            "Requirement already satisfied: importlib-metadata in /usr/local/lib/python3.7/dist-packages (from jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (4.11.0)\n",
            "Requirement already satisfied: zipp>=3.1.0 in /usr/local/lib/python3.7/dist-packages (from importlib-resources>=1.4.0->jsonschema!=2.5.0,>=2.4->nbformat>=4.2.0->ipywidgets->jupyter==1.0.0->d2l) (3.7.0)\n",
            "Requirement already satisfied: Send2Trash in /usr/local/lib/python3.7/dist-packages (from notebook->jupyter==1.0.0->d2l) (1.8.0)\n",
            "Requirement already satisfied: terminado>=0.8.1 in /usr/local/lib/python3.7/dist-packages (from notebook->jupyter==1.0.0->d2l) (0.13.1)\n",
            "Requirement already satisfied: jinja2 in /usr/local/lib/python3.7/dist-packages (from notebook->jupyter==1.0.0->d2l) (2.11.3)\n",
            "Requirement already satisfied: pyzmq>=13 in /usr/local/lib/python3.7/dist-packages (from jupyter-client->ipykernel->jupyter==1.0.0->d2l) (22.3.0)\n",
            "Requirement already satisfied: ptyprocess in /usr/local/lib/python3.7/dist-packages (from terminado>=0.8.1->notebook->jupyter==1.0.0->d2l) (0.7.0)\n",
            "Requirement already satisfied: MarkupSafe>=0.23 in /usr/local/lib/python3.7/dist-packages (from jinja2->notebook->jupyter==1.0.0->d2l) (2.0.1)\n",
            "Requirement already satisfied: testpath in /usr/local/lib/python3.7/dist-packages (from nbconvert->jupyter==1.0.0->d2l) (0.5.0)\n",
            "Requirement already satisfied: defusedxml in /usr/local/lib/python3.7/dist-packages (from nbconvert->jupyter==1.0.0->d2l) (0.7.1)\n",
            "Requirement already satisfied: pandocfilters>=1.4.1 in /usr/local/lib/python3.7/dist-packages (from nbconvert->jupyter==1.0.0->d2l) (1.5.0)\n",
            "Requirement already satisfied: bleach in /usr/local/lib/python3.7/dist-packages (from nbconvert->jupyter==1.0.0->d2l) (4.1.0)\n",
            "Requirement already satisfied: entrypoints>=0.2.2 in /usr/local/lib/python3.7/dist-packages (from nbconvert->jupyter==1.0.0->d2l) (0.4)\n",
            "Requirement already satisfied: mistune<2,>=0.8.1 in /usr/local/lib/python3.7/dist-packages (from nbconvert->jupyter==1.0.0->d2l) (0.8.4)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.7/dist-packages (from bleach->nbconvert->jupyter==1.0.0->d2l) (21.3)\n",
            "Requirement already satisfied: webencodings in /usr/local/lib/python3.7/dist-packages (from bleach->nbconvert->jupyter==1.0.0->d2l) (0.5.1)\n",
            "Requirement already satisfied: qtpy in /usr/local/lib/python3.7/dist-packages (from qtconsole->jupyter==1.0.0->d2l) (2.0.1)\n",
            "Installing collected packages: numpy, requests, pandas, matplotlib, d2l\n",
            "  Attempting uninstall: numpy\n",
            "    Found existing installation: numpy 1.21.5\n",
            "    Uninstalling numpy-1.21.5:\n",
            "      Successfully uninstalled numpy-1.21.5\n",
            "  Attempting uninstall: requests\n",
            "    Found existing installation: requests 2.23.0\n",
            "    Uninstalling requests-2.23.0:\n",
            "      Successfully uninstalled requests-2.23.0\n",
            "  Attempting uninstall: pandas\n",
            "    Found existing installation: pandas 1.3.5\n",
            "    Uninstalling pandas-1.3.5:\n",
            "      Successfully uninstalled pandas-1.3.5\n",
            "  Attempting uninstall: matplotlib\n",
            "    Found existing installation: matplotlib 3.2.2\n",
            "    Uninstalling matplotlib-3.2.2:\n",
            "      Successfully uninstalled matplotlib-3.2.2\n",
            "\u001b[31mERROR: pip's dependency resolver does not currently take into account all the packages that are installed. This behaviour is the source of the following dependency conflicts.\n",
            "tensorflow 2.8.0 requires tf-estimator-nightly==2.8.0.dev2021122109, which is not installed.\n",
            "tensorflow 2.8.0 requires numpy>=1.20, but you have numpy 1.18.5 which is incompatible.\n",
            "tables 3.7.0 requires numpy>=1.19.0, but you have numpy 1.18.5 which is incompatible.\n",
            "google-colab 1.0.0 requires requests~=2.23.0, but you have requests 2.25.1 which is incompatible.\n",
            "datascience 0.10.6 requires folium==0.2.1, but you have folium 0.8.3 which is incompatible.\n",
            "albumentations 0.1.12 requires imgaug<0.2.7,>=0.2.5, but you have imgaug 0.2.9 which is incompatible.\u001b[0m\n",
            "Successfully installed d2l-0.17.3 matplotlib-3.3.3 numpy-1.18.5 pandas-1.2.2 requests-2.25.1\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.colab-display-data+json": {
              "pip_warning": {
                "packages": [
                  "matplotlib",
                  "mpl_toolkits",
                  "numpy",
                  "pandas"
                ]
              }
            }
          },
          "metadata": {}
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Downloading ../data/snli_1.0.zip from https://nlp.stanford.edu/projects/snli/snli_1.0.zip...\n"
          ]
        }
      ],
      "source": [
        "!pip install d2l\n",
        "import os\n",
        "import re\n",
        "import torch\n",
        "from torch import nn\n",
        "from d2l import torch as d2l \n",
        "\n",
        "d2l.DATA_HUB['SNLI'] = (\n",
        "    'https://nlp.stanford.edu/projects/snli/snli_1.0.zip',\n",
        "    '9fcde07509c7e87ec61c640c1b2753d9041758e4')\n",
        "\n",
        "data_dir = d2l.download_extract('SNLI')"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "def read_snli(data_dir, is_train):\n",
        "    \"\"\"将SNLI数据集解析为前提、假设和标签\"\"\"\n",
        "    def extract_text(s):\n",
        "        # 删除我们不会使用的信息\n",
        "        s = re.sub('\\\\(', '', s)\n",
        "        s = re.sub('\\\\)', '', s)\n",
        "        # 用一个空格替换两个或多个连续的空格\n",
        "        s = re.sub('\\\\s{2,}', ' ', s)\n",
        "        return s.strip()#移除字符串头尾空格\n",
        "    label_set = {'entailment': 0, 'contradiction': 1, 'neutral': 2}\n",
        "    file_name = os.path.join(data_dir, 'snli_1.0_train.txt'\n",
        "                             if is_train else 'snli_1.0_test.txt')\n",
        "    with open(file_name, 'r') as f:\n",
        "        rows = [row.split('\\t') for row in f.readlines()[1:]]\n",
        "    premises = [extract_text(row[1]) for row in rows if row[0] in label_set]\n",
        "    hypotheses = [extract_text(row[2]) for row in rows if row[0] \\\n",
        "                in label_set]\n",
        "    labels = [label_set[row[0]] for row in rows if row[0] in label_set]\n",
        "    return premises, hypotheses, labels"
      ],
      "metadata": {
        "id": "Z6F1FiVa3BNM"
      },
      "execution_count": 2,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_data = read_snli(data_dir, is_train=True)\n",
        "for x0, x1, y in zip(train_data[0][:3], train_data[1][:3], train_data[2][:3]):\n",
        "    print('premise:', x0)\n",
        "    print('hypothesis:', x1)\n",
        "    print('label:', y)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ipxvCjjv_4G4",
        "outputId": "c103e66d-fe73-4681-f95f-21f07f42d772"
      },
      "execution_count": 4,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "premise: A person on a horse jumps over a broken down airplane .\n",
            "hypothesis: A person is training his horse for a competition .\n",
            "label: 2\n",
            "premise: A person on a horse jumps over a broken down airplane .\n",
            "hypothesis: A person is at a diner , ordering an omelette .\n",
            "label: 1\n",
            "premise: A person on a horse jumps over a broken down airplane .\n",
            "hypothesis: A person is outdoors , on a horse .\n",
            "label: 0\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "test_data = read_snli(data_dir, is_train=False)\n",
        "for data in [train_data, test_data]:\n",
        "    print([[row for row in data[2]].count(i) for i in range(3)])"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "hWHzhVCIA9EU",
        "outputId": "b0073b1e-deb4-4c4a-c08a-37f590b678ff"
      },
      "execution_count": 5,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "[183416, 183187, 182764]\n",
            "[3368, 3237, 3219]\n"
          ]
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "class SNLIDataset(torch.utils.data.Dataset):\n",
        "    \"\"\"用于加载SNLI数据集的自定义数据集\"\"\"\n",
        "    def __init__(self, dataset, num_steps, vocab=None):\n",
        "        self.num_steps = num_steps\n",
        "        all_premise_tokens = d2l.tokenize(dataset[0])\n",
        "        all_hypothesis_tokens = d2l.tokenize(dataset[1])\n",
        "        if vocab is None:#用bert时一定要使用bert预训练时的vocab,没有的时候自己构建\n",
        "            self.vocab = d2l.Vocab(all_premise_tokens + \\\n",
        "                all_hypothesis_tokens, min_freq=5, reserved_tokens=['<pad>'])\n",
        "        else:\n",
        "            self.vocab = vocab\n",
        "        self.premises = self._pad(all_premise_tokens)\n",
        "        self.hypotheses = self._pad(all_hypothesis_tokens)\n",
        "        self.labels = torch.tensor(dataset[2])\n",
        "        print('read ' + str(len(self.premises)) + ' examples')\n",
        "\n",
        "    def _pad(self, lines):\n",
        "        return torch.tensor([d2l.truncate_pad(\n",
        "            self.vocab[line], self.num_steps, self.vocab['<pad>'])\n",
        "                         for line in lines])\n",
        "\n",
        "    def __getitem__(self, idx):\n",
        "        return (self.premises[idx], self.hypotheses[idx]), self.labels[idx]\n",
        "\n",
        "    def __len__(self):\n",
        "        return len(self.premises)"
      ],
      "metadata": {
        "id": "L3NRAar_BAi5"
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "def load_data_snli(batch_size, num_steps=50):\n",
        "    \"\"\"下载SNLI数据集并返回数据迭代器和词表\"\"\"\n",
        "    num_workers = d2l.get_dataloader_workers()\n",
        "    data_dir = d2l.download_extract('SNLI')\n",
        "    train_data = read_snli(data_dir, True)\n",
        "    test_data = read_snli(data_dir, False)\n",
        "    train_set = SNLIDataset(train_data, num_steps)\n",
        "    test_set = SNLIDataset(test_data, num_steps, train_set.vocab)\n",
        "    train_iter = torch.utils.data.DataLoader(train_set, batch_size,\n",
        "                                             shuffle=True,\n",
        "                                             num_workers=num_workers)\n",
        "    test_iter = torch.utils.data.DataLoader(test_set, batch_size,\n",
        "                                            shuffle=False,\n",
        "                                            num_workers=num_workers)\n",
        "    return train_iter, test_iter, train_set.vocab"
      ],
      "metadata": {
        "id": "TnNpL6veG6aB"
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "train_iter, test_iter, vocab = load_data_snli(128, 50)\n",
        "len(vocab)"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9hAZUEk1HCT6",
        "outputId": "e5d38db0-9d5e-439e-fbcd-b845904e23de"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "read 549367 examples\n",
            "read 9824 examples\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py:481: UserWarning: This DataLoader will create 4 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
            "  cpuset_checked))\n"
          ]
        },
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "18678"
            ]
          },
          "metadata": {},
          "execution_count": 8
        }
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "for X, Y in train_iter:\n",
        "    print(X[0].shape)\n",
        "    print(X[1].shape)\n",
        "    print(Y.shape)\n",
        "    break"
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "5OqGACJTHEcM",
        "outputId": "29221964-b25f-4596-900c-e64a9920db50"
      },
      "execution_count": 9,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stderr",
          "text": [
            "/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py:481: UserWarning: This DataLoader will create 4 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.\n",
            "  cpuset_checked))\n"
          ]
        },
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "torch.Size([128, 50])\n",
            "torch.Size([128, 50])\n",
            "torch.Size([128])\n"
          ]
        }
      ]
    }
  ]
}